import re
import os
import shutil
import tensorflow as tf
from google.protobuf import text_format
from object_detection import exporter
from object_detection.protos import pipeline_pb2

from evaluator.model_evaluator import ModelEvaluator, CkptEvaluator


def get_train_info(train_log):
    ret = {}
    partten = r'global step (\d+): loss = (\d.\d*)'
    ret_lst = re.findall(partten, train_log)
    for step_loss in ret_lst:
        ret[float(step_loss[0])] = float(step_loss[1])
    return ret


def eval_from_chekckpoint(model_dir, config_path, label_map_path, score_threshold):
    if not os.path.exists(model_dir):
        raise FileNotFoundError(model_dir + " Not Found")
    if not os.path.exists(config_path):
        raise FileNotFoundError(config_path + " Not Found")
    latest_checkpoint = tf.train.latest_checkpoint(model_dir)
    if latest_checkpoint is None:
        return {}
    step = int(latest_checkpoint.split('-')[-1])
    evaluator = CkptEvaluator(model_dir, config_path, label_map_path)
    precision, recall = evaluator.get_precison_and_recall(score_threshold)
    mAP = evaluator.get_mAP()
    return {"step": step, "precision": precision, "recall": recall, "mAP": mAP}


def eval_from_pb(pb_file_path, test_images_dir, label_map_path, score_threshold):
    """

    :param pb_file_path:
    :param test_images_dir:
    :param label_map_path:
    :param score_threshold:
    :return:
    :raise: FileNotFoundError if pb_file_path or test_images_dir or label_map_path not exits.
    """
    evaluator = ModelEvaluator(pb_file_path, label_map_path, test_images_dir)
    precision, recall = evaluator.get_precison_and_recall(score_threshold)
    mAP = evaluator.get_mAP()

    return {"precision": precision, "recall": recall, "mAP": mAP}


def export_pb_file(checkpoint_prefix, config_path, output_dir):
    """

    :param checkpoint_prefix:
    :param config_path:
    :param output_dir:
    :return:
    :raise: FileNotFoundError
    """
    if not os.path.exists(checkpoint_prefix+".index"):
        raise FileNotFoundError(checkpoint_prefix + ".index Not Found")
    if not os.path.exists(config_path):
        raise FileNotFoundError(config_path + " Not Found")
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    input_type = 'image_tensor'
    input_shape = None
    config_override = ''
    pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
    with tf.gfile.GFile(config_path, 'r') as f:
        text_format.Merge(f.read(), pipeline_config)
    text_format.Merge(config_override, pipeline_config)
    if input_shape:
        input_shape = [
            int(dim) if dim != '-1' else None
            for dim in input_shape.split(',')
        ]
    else:
        input_shape = None
    exporter.export_inference_graph(input_type, pipeline_config,
                                    checkpoint_prefix,
                                    output_dir, input_shape)
    for path in os.listdir(output_dir):
        full_path = os.path.join(output_dir, path)
        if os.path.isdir(full_path):
            shutil.rmtree(full_path)
        elif path.split('.')[-1] != 'pb':
            os.remove(full_path)

